import io

import numpy as np
from PIL import Image
from matplotlib import pyplot as plt


def visualize(npy_file, n, save_path, show=False):
    assert isinstance(n, int)
    loaded_arr = np.load(npy_file)
    for i, x in enumerate(loaded_arr):
        if i > n != -1:
            break
        plt.imshow(x, cmap='viridis', interpolation='NEAREST', alpha=0.5)
    plt.savefig(save_path)
    if show:
        plt.show()
if __name__ == '__main__':
    visualize(npy_file='visual/yolov8-n-org/test12/13/stage21_C2f_features.npy',
              n=-1,
              save_path='visual/yolov8-n-org/test12/i-3.png')

# # 加载RGB图像
# image_path = './images/6.png'
# image = Image.open(image_path)
#
# # 将图像转换为PyTorch张量并进行必要的预处理
# transform = transforms.Compose([transforms.Resize((224, 224)),
#                                 transforms.ToTensor()])
# input_image = transform(image).unsqueeze(0)  # 添加一个批次维度
#
# # 创建一个卷积层，输出多个特征图
# conv_layer = nn.Conv2d(in_channels=3, out_channels=5, kernel_size=3, padding=1)
#
# # 进行卷积操作
# output_feature_maps = conv_layer(input_image)
#
# # 获取每个特征图的数据
# feature_map_data_list = [output_feature_maps[0, i].detach().numpy() for i in range(output_feature_maps.shape[1])]
#
# # 可视化每个特征图的热力图
# plt.figure(figsize=(12, 6))
# for i, feature_map_data in enumerate(feature_map_data_list):
#     plt.subplot(1, 5, i + 1)
#     plt.imshow(feature_map_data, cmap="viridis")
#     plt.title(f"Feature Map {i + 1}")
#     plt.axis('off')
# plt.show()
